package class03;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Queue;

// 题目描述：
// 给定一个二叉树（具有根结点 root），一个目标结点 target ，和一个整数值 K 。
// 返回到目标结点 target 距离为 K 的所有结点的值的列表。答案可以以任何顺序返回。
//
// 示例：
// 输入：root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, K = 2
// 输出：[7,4,1]
// 解释：所求结点为与目标结点（值为 5）距离为 2 的结点，值分别为 7，4，以及 1
//
// 解题思路：
// 1. 首先建立每个节点与其父节点的映射关系，因为二叉树中节点只有指向子节点的指针，没有指向父节点的指针
// 2. 从target节点开始进行BFS搜索，搜索深度为K
// 3. 在BFS过程中，每个节点可以向三个方向扩展：左子节点、右子节点和父节点
// 4. 使用visited集合避免重复访问节点
// 5. 当到达第K层时，将该层的所有节点加入结果列表
public class Code08_DistanceKNodes {

	public static class Node {
		public int value;
		public Node left;
		public Node right;

		public Node(int v) {
			value = v;
		}
	}

	// 寻找二叉树中距离目标节点为K的所有节点
	public static List<Node> distanceKNodes(Node root, Node target, int K) {
		// 创建一个HashMap用于存储每个节点与其父节点的映射关系
		HashMap<Node, Node> parents = new HashMap<>();
		parents.put(root, null);
		// 建立所有节点的父节点映射关系
		createParentMap(root, parents);
		// 使用队列进行BFS搜索
		Queue<Node> queue = new LinkedList<>();
		// 使用HashSet记录已访问的节点，防止重复访问
		HashSet<Node> visited = new HashSet<>();
		// 将目标节点加入队列和已访问集合
		queue.offer(target);
		visited.add(target);
		// 当前所在的层 level
		int curLevel = 0;
		// 存储结果的列表
		List<Node> ans = new ArrayList<>();
		// 开始BFS搜索
		while (!queue.isEmpty()) {
			// 获取当前层的节点数量
			int size = queue.size();
			// 处理当前层的所有节点
			while (size-- > 0) {
				// 取出队列中的节点
				Node cur = queue.poll();
				// 如果当前层就是第K层，则将节点加入结果列表
				if (curLevel == K) {
					ans.add(cur);
				}
				// 向左子节点扩展（如果左子节点存在且未被访问过）
				if (cur.left != null && !visited.contains(cur.left)) {
					visited.add(cur.left);
					queue.offer(cur.left);
				}
				// 向右子节点扩展（如果右子节点存在且未被访问过）
				if (cur.right != null && !visited.contains(cur.right)) {
					visited.add(cur.right);
					queue.offer(cur.right);
				}
				// 向父节点扩展（如果父节点存在且未被访问过）
				if (parents.get(cur) != null && !visited.contains(parents.get(cur))) {
					visited.add(parents.get(cur));
					queue.offer(parents.get(cur));
				}
			}
			// 处理完当前层后，层数加1
			curLevel++;
			// 如果层数已经超过K，提前结束搜索
			if (curLevel > K) {
				break;
			}
		}
		// 返回结果列表
		return ans;
	}

	// 递归建立每个节点与其父节点的映射关系
	public static void createParentMap(Node cur, HashMap<Node, Node> parents) {
		// 如果当前节点为空，直接返回
		if (cur == null) {
			return;
		}
		// 如果左子节点存在，建立左子节点与当前节点的父子关系映射
		if (cur.left != null) {
			parents.put(cur.left, cur);
			// 递归处理左子树
			createParentMap(cur.left, parents);
		}
		// 如果右子节点存在，建立右子节点与当前节点的父子关系映射
		if (cur.right != null) {
			parents.put(cur.right, cur);
			// 递归处理右子树
			createParentMap(cur.right, parents);
		}
	}

	public static void main(String[] args) {
		// 构造测试用的二叉树
		Node n0 = new Node(0);
		Node n1 = new Node(1);
		Node n2 = new Node(2);
		Node n3 = new Node(3);
		Node n4 = new Node(4);
		Node n5 = new Node(5);
		Node n6 = new Node(6);
		Node n7 = new Node(7);
		Node n8 = new Node(8);

		// 建立节点之间的连接关系
		n3.left = n5;
		n3.right = n1;
		n5.left = n6;
		n5.right = n2;
		n1.left = n0;
		n1.right = n8;
		n2.left = n7;
		n2.right = n4;

		// 设置根节点和目标节点
		Node root = n3;
		Node target = n5;
		int K = 2;

		// 调用方法获取距离为K的节点列表
		List<Node> ans = distanceKNodes(root, target, K);
		// 输出结果
		for (Node o1 : ans) {
			System.out.println(o1.value);
		}

	}

}
